#include "hzpch.h"
#include "OpenGLShader.h"
#include "SourceCode/Renderer/Shader.h"
#include "SourceCode/Shaderc/shaderc.hpp"

#include <glad/glad.h>
#include <glm/gtc/type_ptr.hpp>

#include <spirv_cross.hpp>
#include <spirv_glsl.hpp>

namespace Hazal
{
	namespace Utils
	{
		static const char* GetCacheDirectory()
		{
			return "../../Assets/Cache/Shader/OpenGL";
		}

		static void CreateCacheDirectoryIfNeeded()
		{
			std::string strCacheDir = GetCacheDirectory();
			if (!std::filesystem::exists(strCacheDir))
			{
				std::filesystem::create_directories(strCacheDir);
			}
		}

		static const char* GLShaderStageCachedOpenGLFileExtension(uint32_t stage)
		{
			switch (stage)
			{
			case GL_VERTEX_SHADER:
				return ".cached_opengl.vert";
			case GL_FRAGMENT_SHADER:
				return ".cached_opengl.frag";
			}

			HZ_CORE_ASSERT(false, "Utils::GLShaderStageCachedOpenGLFileExtension Invalid shader type!");

			return "";
		}

		static const char* GLShaderStageCachedVulkanFileExtension(uint32_t stage)
		{
			switch (stage)
			{
			case GL_VERTEX_SHADER:
				return ".cached_vulkan.vert";
			case GL_FRAGMENT_SHADER:
				return ".cached_vulkan.frag";
			}

			HZ_CORE_ASSERT(false, "Utils::GLShaderStageCachedOpenGLFileExtension Invalid shader type!");

			return "";
		}

		static shaderc_shader_kind GlShaderStageToShaderc(GLenum type)
		{
			switch (type)
			{
			case GL_VERTEX_SHADER:
				return shaderc_vertex_shader;
			case GL_FRAGMENT_SHADER:
				return shaderc_fragment_shader;
			default:
				break;
			}

			HZ_CORE_ASSERT(false, "Unknown gl shader type!");
			return (shaderc_shader_kind)0;
		}
	}

	GLenum ShaderTypeFromString(const std::string& strType)
	{
		if ("vertex" == strType)
		{
			return GL_VERTEX_SHADER;
		}
		
		if ("fragment" == strType || "pixel" == strType)
		{
			return GL_FRAGMENT_SHADER;
		}
		
		HZ_CORE_ASSERT(false, "Unknown shader type {0}", strType);
		return GL_NONE;
	}

	COpenGLShader::COpenGLShader(const std::string& strShaderName, const std::string& strVertexSrc, const std::string& strFragmentSrc)
		:m_strShaderName(strShaderName)
	{
		std::unordered_map<GLenum, std::string> mapShaderSrcs;
		mapShaderSrcs[GL_VERTEX_SHADER] = strVertexSrc;
		mapShaderSrcs[GL_FRAGMENT_SHADER] = strFragmentSrc;

		{
			/// TODO 后续补充Timer对象, 计算每次编译所需要的时间
			CompileOrGetVulkanBinaries(mapShaderSrcs);
			CompileOrGetOpenGLBinaries();
			CreateProgram();
		}
	}

	COpenGLShader::COpenGLShader(const std::string& strShaderPath)
		:m_strFilePath(strShaderPath)
	{
		Utils::CreateCacheDirectoryIfNeeded();

		std::string strShaderSrc = ReadFile(strShaderPath);
		auto mapShaderSrcs = PreProcess(strShaderSrc);

		{
			/// TODO 后续补充Timer对象, 计算每次编译所需要的时间
			CompileOrGetVulkanBinaries(mapShaderSrcs);
			CompileOrGetOpenGLBinaries();
			CreateProgram();
		}

		/// Parse shader name
		auto strLastSlash = strShaderPath.find_last_of("/\\");
		strLastSlash = (std::string::npos == strLastSlash) ? 0 : strLastSlash + 1;

		auto strLastDot = strShaderPath.rfind('.');
		strLastDot = (std::string::npos == strLastDot) ? strShaderPath.size() : strLastDot;

		auto nCount = strLastDot - strLastSlash;
		m_strShaderName = strShaderPath.substr(strLastSlash, nCount);
	}

	COpenGLShader::~COpenGLShader()
	{
		glDeleteProgram(m_unRendererID);
	}

	void COpenGLShader::Bind() const
	{
		glUseProgram(m_unRendererID);
	}

	void COpenGLShader::Unbind() const
	{
		glUseProgram(0);
	}

	const std::string& COpenGLShader::GetName() const
	{
		return m_strShaderName;
	}

	void COpenGLShader::SetMat4(const std::string& strUniform, const glm::mat4& mat4)
	{
		UploadUniformMat4(strUniform, mat4);
	}

	void COpenGLShader::SetFloat4(const std::string& strUniform, const glm::vec4& vec4)
	{
		UploadUniformFloat4(strUniform, vec4);
	}

	void COpenGLShader::SetFloat3(const std::string& strUniform, const glm::vec3& vec3)
	{
		UploadUniformFloat3(strUniform, vec3);
	}

	void COpenGLShader::SetFloat(const std::string& strUniform, float fValue)
	{
		UploadUniformFloat(strUniform, fValue);
	}

	void COpenGLShader::SetInt(const std::string& strUniform, int nVal)
	{
		UploadUniformInt(strUniform, nVal);
	}

	void COpenGLShader::SetIntArray(const std::string& strUniform, const int* nVals, uint32_t unArySize)
	{
		UploadUniformIntArray(strUniform, nVals, unArySize);
	}

#pragma region (Override)OpenGL uniform transform operation
	void COpenGLShader::UploadUniformMat4(const std::string& strUniform, const glm::mat4& matrix)
	{
		GLint nLocation = glGetUniformLocation(m_unRendererID, strUniform.c_str());
		glUniformMatrix4fv(nLocation, 1, GL_FALSE, glm::value_ptr(matrix));
	}

	void COpenGLShader::UploadUniformFloat4(const std::string& strUniform, const glm::vec4& vec)
	{
		GLint nLocation = glGetUniformLocation(m_unRendererID, strUniform.c_str());
		glUniform4f(nLocation, vec.x, vec.y, vec.z, vec.w);
	}

	void COpenGLShader::UploadUniformFloat3(const std::string& strUniform, const glm::vec3& vec)
	{
		GLint nLocation = glGetUniformLocation(m_unRendererID, strUniform.c_str());
		glUniform3f(nLocation, vec.x, vec.y, vec.z);
	}

	void COpenGLShader::UploadUniformInt(const std::string& strUniform, const int nSlot)
	{
		GLint nLocation = glGetUniformLocation(m_unRendererID, strUniform.c_str());
		glUniform1i(nLocation, nSlot);
	}

	void COpenGLShader::UploadUniformIntArray(const std::string& strUniform, const int* nVals, uint32_t unArySize)
	{
		GLint nLocation = glGetUniformLocation(m_unRendererID, strUniform.c_str());
		glUniform1iv(nLocation, unArySize, nVals);
	}

	void COpenGLShader::UploadUniformFloat(const std::string& strUniform, float fValue)
	{
		GLint nLocation = glGetUniformLocation(m_unRendererID, strUniform.c_str());
		glUniform1f(nLocation, fValue);
	}
#pragma endregion

	std::string COpenGLShader::ReadFile(const std::string& strFilePath)
	{
		std::string strResult;
		
		std::ifstream in(strFilePath, std::ios::in, std::ios::binary);
		if (in.is_open())
		{
			in.seekg(0, std::ios::end);
			strResult.resize(in.tellg());
			in.seekg(0, std::ios::beg);
			in.read(&strResult[0], strResult.size());
			in.close();
		}
		else
		{
			LOG_CORE_ERROR("Could not open file '{0}'", strFilePath);
		}

		return strResult;
	}

	std::unordered_map<GLenum, std::string> COpenGLShader::PreProcess(const std::string& strShaderSrc)
	{
		std::unordered_map<GLenum, std::string> mapShaderSrcs;
		
		const char* cTypeToken = "#type";
		size_t unTokenSize = strlen(cTypeToken);
		size_t unPos = strShaderSrc.find(cTypeToken, 0);

		while (unPos != std::string::npos)
		{ 
			size_t unEol = strShaderSrc.find_first_of("\r\n", unPos);
			HZ_CORE_ASSERT(unEol != std::string::npos, "GLSL shader fiel parse failed: Syntax error!");
			size_t unBegin = unPos + unTokenSize + 1;
			
			std::string strType = strShaderSrc.substr(unBegin, unEol - unBegin);
			HZ_CORE_ASSERT("vertex" == strType || "fragment" == strType || "pixel", "Invalid shader type specified!");

			size_t unNextLinePos = strShaderSrc.find_first_not_of("\r\n", unEol);
			unPos = strShaderSrc.find(cTypeToken, unNextLinePos);

			mapShaderSrcs[ShaderTypeFromString(strType)] = strShaderSrc.substr(unNextLinePos,
				unPos - (std::string::npos == unNextLinePos ? strShaderSrc.size() - 1 : unNextLinePos));
		}

		return mapShaderSrcs;
	}

#if 0 编译由shaderc完成, 此处暂时保留
	void COpenGLShader::Compile(const std::unordered_map<GLenum, std::string>& mapShaderSrc)
	{
		GLuint unProgram = glCreateProgram();
		
		HZ_CORE_ASSERT(mapShaderSrc.size() <= 2, "Parsed shader type more than 2 types!");
		std::array<GLuint, 2> aryGLShaderIDS;
		GLuint unShaderIdx = 0;

		for (auto& element : mapShaderSrc)
		{
			GLenum eType = element.first;
			const std::string& strShaderSrc = element.second;

			GLuint unShader = glCreateShader(eType);

			const GLchar* pSource = strShaderSrc.c_str();
			glShaderSource(unShader, 1, &pSource, 0);
			glCompileShader(unShader);

			GLint nIsCompiled = 0;
			glGetShaderiv(unShader, GL_COMPILE_STATUS, &nIsCompiled);
			if (GL_FALSE == nIsCompiled)
			{
				GLint nMaxLength = 0;
				glGetShaderiv(unShader, GL_INFO_LOG_LENGTH, &nMaxLength);

				/// Output error information
				std::vector<GLchar> infoLog(nMaxLength);
				glGetShaderInfoLog(unShader, nMaxLength, &nMaxLength, &infoLog[0]);

				glDeleteShader(unShader);

				LOG_CORE_ERROR(infoLog.data());
				LOG_CORE_ERROR("Shader {0} compilation failure!", eType);

				return;
			}

			glAttachShader(unProgram, unShader);

			aryGLShaderIDS[unShaderIdx++] = unShader;
		}

		/// Link process
		glLinkProgram(unProgram);

		GLint nIsLinked = 0;
		glGetProgramiv(unProgram, GL_LINK_STATUS, &nIsLinked);

		if (GL_FALSE == nIsLinked)
		{
			GLint nMaxLength = 0;
			glGetProgramiv(unProgram, GL_INFO_LOG_LENGTH, &nMaxLength);

			/// Output error information
			std::vector<GLchar> infoLog(nMaxLength);
			glGetProgramInfoLog(unProgram, nMaxLength, &nMaxLength, &infoLog[0]);

			glDeleteProgram(unProgram);

			LOG_CORE_ERROR(infoLog.data());
			HZ_CORE_ASSERT(false, "Shader Link failure!");

			for (auto unShader : aryGLShaderIDS)
			{
				glDeleteShader(unShader);
			}

			return;
		}

		/// Always detach shaders after a successful link.
		for (auto unShader : aryGLShaderIDS)
		{
			glDetachShader(unProgram, unShader);
			glDeleteShader(unShader);
		}

		m_unRendererID = unProgram;
	}
#endif

	void COpenGLShader::CompileOrGetVulkanBinaries(const std::unordered_map<GLenum, std::string>& mapShaderSrc)
	{
		/// shaderc编译
		shaderc::Compiler compiler;
		shaderc::CompileOptions options;
		options.SetTargetEnvironment(shaderc_target_env_vulkan, shaderc_env_version_vulkan_1_2);
		
		const bool bOptimize = true;
		if (bOptimize)
		{
			options.SetOptimizationLevel(shaderc_optimization_level_performance);
		}

		std::filesystem::path cacheDirectory = Utils::GetCacheDirectory();
		
		auto& shaderData = m_mapVulkanSPIRV;
		shaderData.clear();

		for (auto&& [stage, source] : mapShaderSrc)
		{
			std::filesystem::path shaderFilePath = m_strFilePath;
			std::filesystem::path cachedPath = cacheDirectory / (shaderFilePath.filename().string() + Utils::GLShaderStageCachedVulkanFileExtension(stage));

			std::ifstream in(cachedPath, std::ios::in | std::ios::binary);

			const char* cTmp = source.c_str();

			/// 如果存在编译文件, 保存至内存; 否则进行编译
			if (in.is_open())
			{
				in.seekg(0, std::ios::end);
				auto size = in.tellg();
				in.seekg(0, std::ios::beg);

				auto& data = shaderData[stage];
				data.resize(size / sizeof(uint32_t));
				in.read((char*)data.data(), size);

				in.close();
			}
			else
			{
				/// TODO: TEST 片段着色器总是编译失败...
				std::string strTemp = "#version 450 core\n \
					struct VertexOutput \
				{ \
					vec4  Color; \
					vec2  TexCoord; \
					float TilingFactor; \
				}; \
				layout(location = 0) in VertexOutput Input; \
				layout(location = 3) in flat float v_TexIndex; \
				layout(location = 4) in flat int v_EntityID; \
				layout(binding = 1) uniform sampler2D u_Textures[32];  \
				layout(location = 0) out vec4 color; \
				layout(location = 1) out int color2; \
				void main() \
				{ \
					color = texture(u_Textures[int(v_TexIndex)], Input.TexCoord * Input.TilingFactor) * Input.Color; \
					color2 = v_EntityID; \
				}";

				if (stage == GL_VERTEX_SHADER)
				{
					strTemp = source;
				}
				shaderc::SpvCompilationResult module = compiler.CompileGlslToSpv(source, Utils::GlShaderStageToShaderc(stage), m_strFilePath.c_str(), options);
				if (shaderc_compilation_status_success != module.GetCompilationStatus())
				{
					LOG_CORE_ERROR("{0}", module.GetErrorMessage());
					HZ_CORE_ASSERT(false, "CompileOrGetVulkanBinaries runtime error!");

					break;
				}

				shaderData[stage] = std::vector<uint32_t>(module.cbegin(), module.cend());

				std::ofstream out(cachedPath, std::ios::out | std::ios::binary);
				if (out.is_open())
				{
					auto& data = shaderData[stage];
					out.write((char*)data.data(), data.size() * sizeof(uint32_t));
					out.flush();
					out.close();
				}
			}
		}

		for (auto&& [stage, data] : shaderData)
		{
			Reflect(stage, data);
		}

	}

	void COpenGLShader::CompileOrGetOpenGLBinaries()
	{
		auto& shaderData = m_mapOpenGLSPIRV;

		shaderc::Compiler compiler;
		shaderc::CompileOptions options;
		options.SetTargetEnvironment(shaderc_target_env_opengl, shaderc_env_version_opengl_4_5);

		const bool bOptimize = false;
		if (bOptimize)
		{
			options.SetOptimizationLevel(shaderc_optimization_level_performance);
		}

		std::filesystem::path cacheDirectory = Utils::GetCacheDirectory();

		shaderData.clear();
		m_umapOpenGLSourceCode.clear();

		for (auto& [stage, spirv] : m_mapVulkanSPIRV)
		{
			std::filesystem::path shaderFilePath = m_strFilePath;
			std::filesystem::path cachePath = cacheDirectory / (shaderFilePath.filename().string() + Utils::GLShaderStageCachedOpenGLFileExtension(stage));

			std::ifstream in(cachePath, std::ios::in | std::ios::binary);

#if 0	///TEST 暂时屏蔽掉, 不知为何每次启动需要重新编译, 否则opengl链接会报错
			if (in.is_open())
			{
				in.seekg(0, std::ios::end);
				auto size = in.tellg();
				LOG_CORE_DEBUG("Bin shader {0} file size = {1}", stage, (int)size);

				in.seekg(0, std::ios::beg);

				auto& data = shaderData[stage];
				data.resize(size / sizeof(uint32_t));
				in.read((char*)data.data(), size);

				in.close();
			}
			else
#endif
			{
				/// Convert to binarg data to source code.
				spirv_cross::CompilerGLSL compilerGLSL(spirv);
				m_umapOpenGLSourceCode[stage] = compilerGLSL.compile();
				auto& source = m_umapOpenGLSourceCode[stage];

				shaderc::SpvCompilationResult module = compiler.CompileGlslToSpv(source, Utils::GlShaderStageToShaderc(stage), m_strFilePath.c_str(), options);
				if (shaderc_compilation_status_success != module.GetCompilationStatus())
				{
					LOG_CORE_ERROR("{0}", module.GetErrorMessage());
					HZ_CORE_ASSERT(false, "CompileOrGetOpenGLBinaries runtime error!");

					break;
				}

				shaderData[stage] = std::vector<uint32_t>(module.cbegin(), module.cend());

				std::ofstream out(cachePath, std::ios::out | std::ios::binary);
				if (out.is_open())
				{
					auto& data = shaderData[stage];
					out.write((char*)data.data(), data.size() * sizeof(size_t));
					out.flush();
					out.close();
				}
			}
		}

	}

	void COpenGLShader::CreateProgram()
	{
		GLint program = glCreateProgram();

		std::vector<GLuint> vecShaderIds;

		for (auto&& [stage, spirv] : m_mapOpenGLSPIRV)
		{
			GLuint unShaderId = vecShaderIds.emplace_back(glCreateShader(stage));
		
			/// 传递编译好的SPIR_V格式二进制文件
			glShaderBinary(1, &unShaderId, GL_SHADER_BINARY_FORMAT_SPIR_V, spirv.data(), spirv.size() * sizeof(uint32_t));
			glSpecializeShader(unShaderId, "main", 0, nullptr, nullptr);
			
			/// 查看编译状态
			GLint nCompileStat = 0;
			glGetShaderiv(unShaderId, GL_COMPILE_STATUS, &nCompileStat);
			if (GL_FALSE == nCompileStat)
			{
				GLint nLogLength = 0;
				glGetShaderiv(unShaderId, GL_INFO_LOG_LENGTH, &nLogLength);

				std::vector<GLchar> vecLogInfo(nLogLength);
				glGetShaderInfoLog(unShaderId, nLogLength, &nLogLength, vecLogInfo.data());

				std::string strLogInfo = vecLogInfo.data();
				LOG_CORE_ERROR("shader {0} compile error: {1}", unShaderId, strLogInfo);

				glDeleteShader(unShaderId);
				return;
			}

			glAttachShader(program, unShaderId);
		}

		glLinkProgram(program);
		GLint nLinkState;
		glGetProgramiv(program, GL_LINK_STATUS, &nLinkState);
		if (GL_FALSE == nLinkState)
		{
			GLint nMaxLength;
			glGetProgramiv(program, GL_INFO_LOG_LENGTH, &nMaxLength);

			std::vector<GLchar> vecInfo(nMaxLength);
			glGetProgramInfoLog(program, nMaxLength, &nMaxLength, vecInfo.data());

			LOG_CORE_ERROR("Shader link error({0}): {1}", m_strFilePath, vecInfo.data());
			  
			glDeleteProgram(program);

			for (auto unShaderId : vecShaderIds)
			{
				glDeleteShader(unShaderId);
			}
		}

		for (auto unShaderId : vecShaderIds)
		{
			glDetachShader(program, unShaderId);
			glDeleteShader(unShaderId);
		}

		m_unRendererID = program;
	}

	void COpenGLShader::Reflect(GLenum stage, const std::vector<uint32_t>& vecShaderData)
	{
		spirv_cross::Compiler compiler(vecShaderData);
		spirv_cross::ShaderResources resources = compiler.get_shader_resources();

		LOG_CORE_INFO("{0} uniform buffers", resources.uniform_buffers.size());
		LOG_CORE_INFO("{0} resources", resources.sampled_images.size());
	}


}